from mmengine.config import read_base
with read_base():
    from .mgam import *

from mmengine.model.weight_init import PretrainedInit
from mgamdata.models.MedNeXt import MM_MedNext_Encoder, MM_MedNext_Decoder_3D
from mgamdata.mm.mmseg_Dev3D import DiceLoss_3D, EncoderDecoder_3D


# val_evaluator=val_cfg=val_dataloader=multi_dataset_Ts=None

# 神经网络设定
model = dict(
    type = EncoderDecoder_3D,
    backbone = dict(
        type=MM_MedNext_Encoder,
        init_cfg = dict(
            type=PretrainedInit, 
            checkpoint=pretrained_model,
            prefix="backbone.0",
            map_location="cpu"
        ) if pretrained_model is not None else None,
        in_channels=in_channels, # type: ignore
        embed_dims=embed_dims,
        kernel_size=3,
        dim="3d",
        use_checkpoint=use_checkpoint,
        norm_type='layer',
        freeze=True,
    ),
    decode_head = dict(
        type=MM_MedNext_Decoder_3D,
        init_cfg = dict(
            type=PretrainedInit, 
            checkpoint=pretrained_model,
            prefix="backbone.2",
            map_location="cpu"
        ) if pretrained_model is not None else None,
        embed_dims=embed_dims,
        kernel_size=3,
        num_classes=embed_dims,
        out_channels=embed_dims,
        use_checkpoint=use_checkpoint,
        deep_supervision=deep_supervision,
        norm_type='layer',
        ignore_index=0, # 仅对train acc计算有效
        loss_gt_key='gt_sem_seg', # ["gt_sem_seg_one_hot", "gt_sem_seg"]
        loss_decode=dict(
            type=DiceLoss_3D, 
            use_sigmoid=False, 
            ignore_1st_index=False, 
            batch_z=4, 
            # NOTE Severe performance overhead when not being set to None.
            # NOTE Prefer using `ignore_1st_index`.
            # NOTE Invalid Class (Defaults to the last class) has been masked out during preprocess.
            ignore_index=None, 
        ), 
    ),
    test_cfg=dict(
        mode="slide",
        crop_size=size,
        slide_accumulate_device='cpu',
        stride=[i//2 for i in size]
    )
)